#include "stdafx.h"
#include "ps.h"
#include "nt.h"

DWORD NewEnumProcess(std::set<DWORD> &v)
{
	ULONG cbBuffer = 0;
	LPVOID pBuffer = NULL;
	DWORD nCount = 0;

	while (1)
	{
		cbBuffer += 0x20000;
		pBuffer = VirtualAlloc(NULL, cbBuffer, MEM_COMMIT, PAGE_READWRITE);

		if (pBuffer == NULL)
		{
			return 0;
		}

		NTSTATUS Status = NewNtQuerySystemInformation(SystemProcessInformation, pBuffer, cbBuffer, NULL);

		if (NT_SUCCESS(Status))
		{
			break;
		}

		VirtualFree(pBuffer, 0, MEM_RELEASE);

		if (Status != STATUS_INFO_LENGTH_MISMATCH)
		{
			return 0;
		}
	}

	if (pBuffer == NULL)
		return 0;

	PSYSTEM_PROCESS_INFORMATION pInfo = (PSYSTEM_PROCESS_INFORMATION)pBuffer;

	while (pInfo)
	{
		DWORD dwProcessId = (DWORD)pInfo->UniqueProcessId;

		if (dwProcessId > 4)
		{
			v.insert(dwProcessId);
			nCount++;
		}

		if (pInfo->NextEntryOffset == 0)
		{
			break;
		}

		pInfo = (PSYSTEM_PROCESS_INFORMATION)(((PUCHAR)pInfo) + pInfo->NextEntryOffset);
	}

	VirtualFree(pBuffer, 0, MEM_RELEASE);
	return nCount;
}

DWORD NewEnumProcessEx(std::set<DWORD>&v)
{
	ULONG cbBuffer = 0;
	LPVOID pBuffer = NULL;
	DWORD nCount = 0;

	while (1)
	{
		cbBuffer += 0x20000;
		pBuffer = VirtualAlloc(NULL, cbBuffer, MEM_COMMIT, PAGE_READWRITE);

		if (pBuffer == NULL)
		{
			return 0;
		}

		NTSTATUS Status = NewNtQuerySystemInformation(SystemHandleInformation, pBuffer, cbBuffer, NULL);

		if (NT_SUCCESS(Status))
		{
			break;
		}

		VirtualFree(pBuffer, 0, MEM_RELEASE);

		if (Status != STATUS_INFO_LENGTH_MISMATCH)
		{
			return 0;
		}
	}

	if (pBuffer == NULL)
		return 0;

	PSYSTEM_HANDLE_INFORMATION pInfo = (PSYSTEM_HANDLE_INFORMATION)pBuffer;

	for (ULONG i = 0; i < pInfo->NumberOfHandles; ++i)
	{
		DWORD dwProcessId = *(DWORD *)&pInfo->Handles[i].UniqueProcessId;

		if (dwProcessId > 4)
		{
			v.insert(dwProcessId);
			nCount++;
		}
	}

	VirtualFree(pBuffer, 0, MEM_RELEASE);
	return nCount;
}

BOOL GetNtoskrnlImagePath(LPWSTR szImagePath)
{
	ULONG cbBuffer = 0;
	LPVOID pBuffer = NULL;
	BOOL bSuccess = FALSE;

	while (1)
	{
		cbBuffer += 0x20000;
		pBuffer = VirtualAlloc(NULL, cbBuffer, MEM_COMMIT, PAGE_READWRITE);

		if (pBuffer == NULL)
		{
			return FALSE;
		}

		NTSTATUS Status = NewNtQuerySystemInformation(SystemModuleInformation, pBuffer, cbBuffer, NULL);

		if (NT_SUCCESS(Status))
		{
			break;
		}

		VirtualFree(pBuffer, 0, MEM_RELEASE);

		if (Status != STATUS_INFO_LENGTH_MISMATCH)
		{
			return FALSE;
		}
	}

	if (pBuffer == NULL)
		return FALSE;

	PSYSTEM_MODULE_INFORMATION pInfo = (PSYSTEM_MODULE_INFORMATION)pBuffer;
	for (size_t i = 0; i < pInfo->Count; ++i)
	{
		PSYSTEM_MODULE_INFORMATION_ENTRY pEntry = &pInfo->Module[i];
		if (pEntry->LoadOrderIndex == 0)
		{
			if (MultiByteToWideChar(CP_ACP, 0, pInfo->Module[i].ImageName, strlen(pInfo->Module[i].ImageName), szImagePath, MAX_PATH) > 0)
			{
				bSuccess = TRUE;
			}
			break;
		}
	}

	VirtualFree(pBuffer, 0, MEM_RELEASE);
	return bSuccess;
}